{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "from netCDF4 import Dataset\n",
    "import numpy as np\n",
    "import numpy.ma as ma\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "from matplotlib.cm import get_cmap\n",
    "from mpl_toolkits.basemap import Basemap\n",
    "import matplotlib.gridspec as gridspec\n",
    "from wrf import getvar, interplevel, to_np, get_basemap, latlon_coords, destagger\n",
    "import glob\n",
    "from matplotlib.colors import LinearSegmentedColormap\n",
    "import cmocean\n",
    "import pyresample\n",
    "import geopy.distance\n",
    "from scipy.ndimage import gaussian_filter1d\n",
    "from scipy import interpolate\n",
    "import os\n",
    "import Nio\n",
    "import datetime as dt\n",
    "import seaborn as sns\n",
    "import pandas as pd\n",
    "import copy\n",
    "import xarray as xr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### !!! User mods !!!\n",
    "#### How many hours before ice edge do you want to start?\n",
    "#### Note: t0_h = 0 means you are starting approx. at ice edge, t0_h = 10 means you are starting 10 h upstream (north) of the ice edge\n",
    "#### Note: t0_h must be an integer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "t0_h = 0\n",
    "if t0_h < 0 or t0_h > 10:\n",
    "    sys.exit('Error: Please set 0 <= t0_h >= 10')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set some things"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Start time is: 2020-03-13 00:00:00\n"
     ]
    }
   ],
   "source": [
    "nhrs = 18 + t0_h + 1 # total number of simulation hours, including t0; default is from 2020-03-12 at 14 UTC to 2020-03-14 at 00 UTC\n",
    "if t0_h == 0:\n",
    "    start_time = '2020-03-13 00:00:00'\n",
    "    start_day = 13\n",
    "    start_hour = 0\n",
    "else:\n",
    "    start_time = '2020-03-12 ' + str(24-t0_h) + ':00:00'\n",
    "    start_day = 12\n",
    "    start_hour = 24-t0_h\n",
    "print ('Start time is: ' + start_time)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get LES domain locations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "LES domain mid point: lat=73.1N, lon=8.7E\n",
      "[ 2.  3.  4.  5.  6.  7.  8.  9. 10. 11. 12. 13. 14. 15. 16. 17. 18.]\n",
      "[76.578 76.005 75.454 74.905 74.356 73.832 73.345 72.884 72.438 71.998\n",
      " 71.559 71.116 70.679 70.265 69.863 69.474 69.141]\n",
      "[ 2.121  3.177  4.29   5.487  6.622  7.603  8.479  9.319 10.145 10.95\n",
      " 11.748 12.519 13.201 13.78  14.425 15.108 15.684]\n",
      "17\n"
     ]
    }
   ],
   "source": [
    "#fname = '/glade/scratch/tjuliano/doe_comble/intercomparison/dephy_forcing/LES_domain_location_28h_18Z_Mar13_2020.txt'\n",
    "#fname = '/glade/scratch/tjuliano/doe_comble/intercomparison/dephy_forcing/LES_domain_location_36h_12Z_Mar13_2020.txt'\n",
    "fname = 'LES_domain_location_era5_36h_12Z_Mar13_2020.txt'\n",
    "\n",
    "les_loc = np.loadtxt(fname,skiprows=1)\n",
    "les_hh = les_loc[:,0]\n",
    "les_hh = les_hh + 6\n",
    "les_hh_idx = np.where(les_hh>=-1.*t0_h)[0]\n",
    "les_hh_refined = les_hh[les_hh_idx]\n",
    "les_hh_refined = les_hh_refined[::2]\n",
    "les_hh_refined = les_hh_refined[2:]\n",
    "les_lat = les_loc[les_hh_idx,1]\n",
    "#les_lat2 = les_lat[::12]\n",
    "les_lat2 = les_lat[::2]\n",
    "les_lat2 = les_lat2[2:]\n",
    "les_lon = les_loc[les_hh_idx,2]\n",
    "#les_lon2 = les_lon[::12]\n",
    "les_lon2 = les_lon[::2]\n",
    "les_lon2 = les_lon2[2:]\n",
    "\n",
    "les_lat_mid = round(np.mean(les_lat),1)\n",
    "les_lon_mid = round(np.mean(les_lon),1)\n",
    "print ('LES domain mid point: ' + 'lat=' + str(les_lat_mid) + 'N, lon=' + str(les_lon_mid) + 'E')\n",
    "\n",
    "#print (les_hh_refined[::12])\n",
    "print (les_hh_refined)\n",
    "print (les_lat2)\n",
    "print (les_lon2)\n",
    "print (len(les_lon2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "savedir = \"figs_new_nz136/\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Read files for plotting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "# precip vars\n",
    "raw_precip_traj = xr.open_dataset('raw_precip_traj_mar13.nc')\n",
    "mean_precip_no_zeros_arr = raw_precip_traj['mean_precip_no_zeros_arr']\n",
    "where_precip = raw_precip_traj['where_precip']\n",
    "\n",
    "# hydromet vars\n",
    "raw_hydromet_traj = xr.open_dataset('raw_hydromet_traj_mar13.nc')\n",
    "hydromet_arr = raw_hydromet_traj['hydromet_arr']\n",
    "\n",
    "# mean state vars\n",
    "raw_mean_var_traj = xr.open_dataset('raw_mean_var_traj_mar13.nc')\n",
    "mean_var_arr = raw_mean_var_traj['mean_var_arr']\n",
    "\n",
    "# theta\n",
    "raw_theta_traj = xr.open_dataset('raw_theta_traj_mar13.nc')\n",
    "theta_arr = raw_theta_traj['theta_arr']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Trajectory subplot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(22,12))\n",
    "grid = plt.GridSpec(11, 3, hspace=0.2, wspace=0.325)\n",
    "main_ax = fig.add_subplot(grid[:, 0])\n",
    "\n",
    "cols = ['lightgreen','lightblue','darkgreen','darkblue']\n",
    "labels = [\"Level 2.5\", \"Level 3\", \"Level 2.5 EDMF\", \"YSU\"]\n",
    "fs = 18\n",
    "xlab_hrs = ['-16','-15','-14','-13','-12','-11','-10','-9','-8','-7','-6','-5','-4','-3','-2','-1','0']\n",
    "xlab_empty = ['','','','','','','','','','','','','','','','','']\n",
    "\n",
    "ref_lon = (les_lon[0]+les_lon[-1])/2.\n",
    "ref_lat = (les_lat[0]+les_lat[-1])/2.\n",
    "lono = 0\n",
    "lato = 5\n",
    "ll_lon = 2\n",
    "ll_lat = 68.25\n",
    "ul_lon = 2\n",
    "ul_lat = 77\n",
    "ur_lon = 22\n",
    "ur_lat = 77\n",
    "lr_lon = 22\n",
    "lr_lat = 69.25\n",
    "# center point of overall domain\n",
    "ref_lat = (ll_lat + ur_lat)/2.\n",
    "#ref_lon = -(abs(ll_lon) + abs(ur_lon))/2.\n",
    "ref_lon = 8.5\n",
    "#m = Basemap(llcrnrlon=les_lon[-1]-lono,llcrnrlat=les_lat[-1]-lato,urcrnrlon=les_lon[0]+lono,urcrnrlat=les_lat[0]+lato,lon_0=ref_lon,lat_0=ref_lat,projection='cass',resolution='i')\n",
    "m = Basemap(llcrnrlon=ll_lon,llcrnrlat=ll_lat,urcrnrlon=ur_lon,urcrnrlat=ur_lat,lon_0=ref_lon,lat_0=ref_lat,projection='cass',resolution='i',ax=main_ax)\n",
    "x, y = m(to_np(les_lon2), to_np(les_lat2))\n",
    "col = np.arange(len(les_hh_refined))\n",
    "m.plot(x,y,c='k',lw=3.,zorder=4)\n",
    "m.scatter(x,y,facecolor='blue',edgecolor='k',s=400,marker='X',zorder=4)\n",
    "m.scatter(x[-13],y[-13],facecolor='lightskyblue',edgecolor='k',s=400,marker='X',zorder=4)\n",
    "m.scatter(x[-7],y[-7],facecolor='lightskyblue',edgecolor='k',s=400,marker='X',zorder=4)\n",
    "m.scatter(x[-1],y[-1],facecolor='lightskyblue',edgecolor='k',s=400,marker='X',zorder=4)\n",
    "for i in np.arange(0,len(x),2):\n",
    "    if i < 8:\n",
    "        xoff = x[6] - x[4]\n",
    "        yoff = y[2] - y[1]\n",
    "        plt.text(x[i]-1.25*xoff,y[i]+0.5*yoff,str(int(les_hh_refined[i])-18),fontsize=18)\n",
    "    else:\n",
    "        xoff = x[3] - x[1]\n",
    "        yoff = y[2] - y[1]\n",
    "        plt.text(x[i]-xoff,y[i]+0.5*yoff,str(int(les_hh_refined[i])-18),fontsize=18)\n",
    "\n",
    "# Add geographic outlines\n",
    "m.drawcoastlines(linewidth=2.5,zorder=3)\n",
    "m.drawparallels(np.arange(60.,94.,2),labels=[True,False,False,False],fontsize=22,linewidth=3.0,zorder=3)\n",
    "m.drawmeridians(np.arange(-40.,70.,4),labels=[False,False,False,True],fontsize=22,linewidth=3.0,zorder=3)\n",
    "m.drawcountries(linewidth=2.0,zorder=3)\n",
    "m.fillcontinents(color='coral',lake_color='coral',zorder=3)\n",
    "\n",
    "ms = 90\n",
    "fs = 18\n",
    "xlab_hrs = ['-16','','-14','','-12','','-10','','-8','','-6','','-4','','-2','','0']\n",
    "\n",
    "### PRECIP RATE\n",
    "xidx = 1\n",
    "dim0 = np.shape(mean_precip_no_zeros_arr)[0]\n",
    "dim2 = np.shape(mean_precip_no_zeros_arr)[2]\n",
    "pfact = 3600. # convert mm/s to mm/h\n",
    "\n",
    "main_ax2 = fig.add_subplot(grid[1:4, xidx])\n",
    "for i in np.arange(dim0):\n",
    "    plt.plot(np.arange(dim2),pfact*mean_precip_no_zeros_arr[i,0,:],c=cols[i],zorder=2)\n",
    "    plt.scatter(np.arange(dim2),pfact*mean_precip_no_zeros_arr[i,0,:],s=ms,facecolor=cols[i],edgecolor='k',lw=1.,label=labels[i],zorder=3)\n",
    "    plt.ylabel('$P_{Rain}$ [mm/h]',fontsize=fs)\n",
    "    plt.xlim(-0.5,16.5)\n",
    "    plt.xticks(np.arange(dim2),xlab_empty,fontsize=fs)\n",
    "    plt.yticks(fontsize=fs)\n",
    "    plt.grid(True)\n",
    "    plt.legend(fontsize=14)\n",
    "\n",
    "main_ax2 = fig.add_subplot(grid[4:7, xidx])\n",
    "for i in np.arange(dim0):\n",
    "    plt.plot(np.arange(dim2),pfact*mean_precip_no_zeros_arr[i,1,:],c=cols[i])\n",
    "    plt.scatter(np.arange(dim2),pfact*mean_precip_no_zeros_arr[i,1,:],s=ms,facecolor=cols[i],edgecolor='k',lw=1.,label=labels[i],zorder=3)\n",
    "    plt.ylabel('$P_{Snow}$ [mm/h]',fontsize=fs)\n",
    "    plt.xlim(-0.5,16.5)\n",
    "    plt.xticks(np.arange(dim2),xlab_empty,fontsize=fs)\n",
    "    plt.yticks(fontsize=fs)\n",
    "    plt.grid(True)\n",
    "    \n",
    "main_ax2 = fig.add_subplot(grid[7:10, xidx])\n",
    "for i in np.arange(dim0):\n",
    "    plt.plot(np.arange(dim2),pfact*mean_precip_no_zeros_arr[i,2,:],c=cols[i])\n",
    "    plt.scatter(np.arange(dim2),pfact*mean_precip_no_zeros_arr[i,2,:],s=ms,facecolor=cols[i],edgecolor='k',lw=1.,label=labels[i],zorder=3)\n",
    "    plt.xlabel('Time before arriving at Andenes (h)',fontsize=fs)\n",
    "    plt.ylabel('$P_{Graupel}$ [mm/h]',fontsize=fs)\n",
    "    plt.xlim(-0.5,16.5)\n",
    "    plt.xticks(np.arange(dim2),xlab_hrs,fontsize=fs)\n",
    "    plt.yticks(fontsize=fs)\n",
    "    plt.grid(True)\n",
    "\n",
    "    \n",
    "### PRECIP FRACTION \n",
    "xidx = 2\n",
    "main_ax2 = fig.add_subplot(grid[1:4, xidx])\n",
    "for i in np.arange(dim0):\n",
    "    plt.plot(np.arange(dim2),where_precip[i,0,:],c=cols[i])\n",
    "    plt.scatter(np.arange(dim2),where_precip[i,0,:],s=ms,facecolor=cols[i],edgecolor='k',lw=1.,label=labels[i],zorder=3)\n",
    "    plt.ylabel('$A_{Rain}$ [%]',fontsize=fs)\n",
    "    plt.xlim(-0.5,16.5)\n",
    "    plt.xticks(np.arange(dim2),xlab_empty,fontsize=fs)\n",
    "    plt.yticks(fontsize=fs)\n",
    "    plt.grid(True)\n",
    "    \n",
    "main_ax3 = fig.add_subplot(grid[4:7, xidx])\n",
    "for i in np.arange(dim0):    \n",
    "    plt.plot(np.arange(dim2),where_precip[i,1,:],c=cols[i])\n",
    "    plt.scatter(np.arange(dim2),where_precip[i,1,:],s=ms,facecolor=cols[i],edgecolor='k',lw=1.,label=labels[i],zorder=3)\n",
    "    plt.ylabel('$A_{Snow}$ [%]',fontsize=fs)\n",
    "    plt.xlim(-0.5,16.5)\n",
    "    plt.xticks(np.arange(dim2),xlab_empty,fontsize=fs)\n",
    "    plt.yticks(fontsize=fs)\n",
    "    plt.grid(True)\n",
    "\n",
    "main_ax4 = fig.add_subplot(grid[7:10, xidx])\n",
    "for i in np.arange(dim0):    \n",
    "    plt.plot(np.arange(dim2),where_precip[i,2,:],c=cols[i])\n",
    "    plt.scatter(np.arange(dim2),where_precip[i,2,:],s=ms,facecolor=cols[i],edgecolor='k',lw=1.,label=labels[i],zorder=3)\n",
    "    plt.xlabel('Time before arriving at Andenes (h)',fontsize=fs)\n",
    "    plt.ylabel('$A_{Graupel}$ [%]',fontsize=fs)\n",
    "    plt.xlim(-0.5,16.5)\n",
    "    plt.xticks(np.arange(dim2),xlab_hrs,fontsize=fs)\n",
    "    plt.yticks(fontsize=fs)\n",
    "    plt.grid(True)\n",
    "\n",
    "plt.savefig(savedir+'traj_tseries.png',dpi=600,bbox_inches='tight')\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FIG 11: PANEL SHOWING VPROFS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "var_lab = ['rthblten','rqvblten','rqcblten','rqiblten','rqniblten','rublten','rvblten','exch_h','exch_m','el_pbl','qke','sh3d','h_diabatic','qv_diabatic','qc_diabatic','ph']\n",
    "# np.empty([len(files_tends),len(calipso_j1),len(var_lab),np.shape(var)[0]])\n",
    "labels = [\"Level 2.5\", \"Level 3\", \"Level 2.5\\nEDMF\", \"YSU\"]\n",
    "\n",
    "z1 = 0\n",
    "z2 = 110\n",
    "lw = 2.5\n",
    "fs = 12\n",
    "fs2 = 14\n",
    "colors = ['lightgreen','lightblue','darkgreen','darkblue']\n",
    "time_arr = [-13,-7,-1]\n",
    "sp_loc = [[1,8,15],[2,9,16],[3,10,17],[4,11,18],[5,12,19],[6,13,20],[7,14,21]]\n",
    "ylim2 = 5500\n",
    "yticks = [0,1000,2000,3000,4000,5000]\n",
    "yticklabs = ['','','','','','']\n",
    "\n",
    "hydromet_arr2 = copy.deepcopy(hydromet_arr)\n",
    "hydromet_arr2 = 1000.*hydromet_arr2 # kg/kg --> g/kg\n",
    "\n",
    "plt.figure(figsize=(15,9))\n",
    "for i in np.arange(4): # sims\n",
    "    for j in np.arange(3): # times\n",
    "        plt.subplot(3,7,sp_loc[0][j])\n",
    "        plt_varx = hydromet_arr2[i,time_arr[j],0,z1:z2]\n",
    "        #plt_vary = mean_var_arr[i,time_arr[j],-1,z1:z2]\n",
    "        plt_vary = theta_arr[i,1,time_arr[j],z1:z2]\n",
    "        plt.plot(plt_varx,plt_vary,c=colors[i],lw=lw,zorder=2) # KH/KM\n",
    "        #plt.scatter(plt_varx[::8],plt_vary[::8],s=60,facecolor=colors[i],edgecolor='k',zorder=2)\n",
    "        plt.xlim(0,2e-2)\n",
    "        if j < 2:\n",
    "            plt.xticks([0,1e-2,2e-2],['','',''])\n",
    "        else:\n",
    "            plt.ticklabel_format(style='sci', axis='x', scilimits=(0,0))\n",
    "            plt.xticks([0,1e-2,2e-2],fontsize=fs)\n",
    "            plt.xlabel('[g kg$^{-1}$]',fontsize=fs)\n",
    "        plt.ylim(0,ylim2)\n",
    "        plt.yticks(yticks,fontsize=fs)\n",
    "        plt.ylabel('Height ASL [m]',fontsize=fs2)\n",
    "        plt.grid(True,zorder=1)\n",
    "        if j == 0:\n",
    "            plt.title('Q$_{c}$',fontsize=fs2)\n",
    "        \n",
    "        \n",
    "        \n",
    "        plt.subplot(3,7,sp_loc[1][j])\n",
    "        plt_varx = hydromet_arr2[i,time_arr[j],1,z1:z2]\n",
    "        #plt_vary = mean_var_arr[i,time_arr[j],-1,z1:z2]\n",
    "        plt_vary = theta_arr[i,1,time_arr[j],z1:z2]\n",
    "        plt.plot(plt_varx,plt_vary,c=colors[i],lw=lw,zorder=2) # KH/KM\n",
    "        #plt.scatter(plt_varx[::8],plt_vary[::8],s=60,facecolor=colors[i],edgecolor='k',zorder=2)\n",
    "        plt.xlim(0,2e-4)\n",
    "        if j < 2:\n",
    "            plt.xticks([0,1e-4,2e-4],['','',''])\n",
    "        else:\n",
    "            plt.ticklabel_format(style='sci', axis='x', scilimits=(0,0))\n",
    "            plt.xticks([0,1e-4,2e-4],fontsize=fs)\n",
    "            plt.xlabel('[g kg$^{-1}$]',fontsize=fs)\n",
    "        plt.ylim(0,ylim2)\n",
    "        plt.yticks(yticks,yticklabs)\n",
    "        plt.grid(True,zorder=1)\n",
    "        if j == 0:\n",
    "            plt.title('Q$_{r}$',fontsize=fs2)\n",
    "        \n",
    "        \n",
    "        \n",
    "        ax3 = plt.subplot(3,7,sp_loc[2][j])\n",
    "        plt_varx = hydromet_arr2[i,time_arr[j],2,z1:z2]\n",
    "        #plt_vary = mean_var_arr[i,time_arr[j],-1,z1:z2]\n",
    "        plt_vary = theta_arr[i,1,time_arr[j],z1:z2]\n",
    "        ax3.plot(plt_varx,plt_vary,c=colors[i],lw=lw,zorder=2) # KH/KM\n",
    "        #plt.scatter(plt_varx[::8],plt_vary[::8],s=60,facecolor=colors[i],edgecolor='k',zorder=2)\n",
    "        ax3.set_xlim(0,3e-2)\n",
    "        if j < 2:\n",
    "            ax3.set_xticks([0,1e-2,2e-2,3e-2],['','','',''])\n",
    "        else:\n",
    "            ax3.ticklabel_format(style='sci', axis='x', scilimits=(0,0))\n",
    "            plt.xticks([0,1e-2,2e-2,3e-2],fontsize=fs)\n",
    "            ax3.set_xlabel('[g kg$^{-1}$]',fontsize=fs)\n",
    "        ax3.set_ylim(0,ylim2)\n",
    "        ax3.set_yticks(yticks,yticklabs)\n",
    "        plt.grid(True,zorder=1)\n",
    "        if j == 0:\n",
    "            plt.title('Q$_{i}$',fontsize=fs2)\n",
    "            \n",
    "        #ax3b = ax3.twiny()\n",
    "        #plt_varx = theta_arr[i,0,time_arr[j],z1:z2] - 273.15\n",
    "        #plt_vary = theta_arr[i,1,time_arr[j],z1:z2]\n",
    "        #ax3b.plot(plt_varx,plt_vary,c='k',lw=lw,zorder=1)\n",
    "        #ax3b.plot([-35,-35],[plt_vary[0],plt_vary[-1]],c='b',lw=lw,ls='--',zorder=1)\n",
    "        #ax3b.set_xlim(-50,0)\n",
    "        \n",
    "        \n",
    "        \n",
    "        plt.subplot(3,7,sp_loc[3][j])\n",
    "        plt_varx = hydromet_arr2[i,time_arr[j],3,z1:z2]\n",
    "        #plt_vary = mean_var_arr[i,time_arr[j],-1,z1:z2]\n",
    "        plt_vary = theta_arr[i,1,time_arr[j],z1:z2]\n",
    "        plt.plot(plt_varx,plt_vary,c=colors[i],lw=lw,zorder=2) # KH/KM\n",
    "        #plt.scatter(plt_varx[::8],plt_vary[::8],s=60,facecolor=colors[i],edgecolor='k',zorder=2)\n",
    "        plt.xlim(0,3.6e-1)\n",
    "        if j < 2:\n",
    "            plt.xticks([0,2e-1],['',''])\n",
    "        else:\n",
    "            plt.ticklabel_format(style='sci', axis='x', scilimits=(0,0))\n",
    "            plt.xticks([0,2e-1],fontsize=fs)\n",
    "            plt.xlabel('[g kg$^{-1}$]',fontsize=fs)\n",
    "        plt.ylim(0,ylim2)\n",
    "        plt.yticks(yticks,yticklabs)\n",
    "        plt.grid(True,zorder=1)\n",
    "        if j == 0:\n",
    "            plt.title('Q$_{s}$',fontsize=fs2)\n",
    "            \n",
    "        \n",
    "        \n",
    "        ax5 = plt.subplot(3,7,sp_loc[4][j])\n",
    "        plt_varx = hydromet_arr2[i,time_arr[j],4,z1:z2]\n",
    "        #plt_vary = mean_var_arr[i,time_arr[j],-1,z1:z2]\n",
    "        plt_vary = theta_arr[i,1,time_arr[j],z1:z2]\n",
    "        plt.plot(plt_varx,plt_vary,c=colors[i],lw=lw,zorder=2,label=labels[i]) # KH/KM\n",
    "        #plt.scatter(plt_varx[::8],plt_vary[::8],s=60,facecolor=colors[i],edgecolor='k',zorder=2)\n",
    "        plt.xlim(0,4e-3)\n",
    "        if j < 2:\n",
    "            plt.xticks([0,2e-3,4e-3],['','',''])\n",
    "        else:\n",
    "            plt.ticklabel_format(style='sci', axis='x', scilimits=(0,0))\n",
    "            plt.xticks([0,2e-3,4e-3],fontsize=fs)\n",
    "            plt.xlabel('[g kg$^{-1}$]',fontsize=fs)\n",
    "        plt.ylim(0,ylim2)\n",
    "        plt.yticks(yticks,yticklabs)\n",
    "        plt.grid(True,zorder=1)\n",
    "        if j == 0:\n",
    "            plt.title('Q$_{g}$',fontsize=fs2)\n",
    "            plt.legend()\n",
    "                    \n",
    "        \n",
    "        \n",
    "        \n",
    "        ax6 = plt.subplot(3,7,sp_loc[5][j])\n",
    "        plt_varx1 = mean_var_arr[i,time_arr[j],7,z1:z2]\n",
    "        plt_varx2 = mean_var_arr[i,time_arr[j],8,z1:z2]\n",
    "        plt_vary = mean_var_arr[i,time_arr[j],-1,z1:z2]\n",
    "        plt.plot(plt_varx1,plt_vary,c=colors[i],lw=lw,zorder=2)\n",
    "        plt.plot(plt_varx2,plt_vary,c=colors[i],lw=lw,ls='--',zorder=2)\n",
    "        #plt_vary = mean_var_arr[i,time_arr[j],-1,z1:z2]\n",
    "        plt_vary = theta_arr[i,1,time_arr[j],z1:z2]\n",
    "        #plt.scatter(plt_varx[::8],plt_vary[::8],s=60,facecolor=colors[i],edgecolor='k',zorder=2)\n",
    "        plt.xlim(0,4e-3)\n",
    "        if j < 2:\n",
    "            plt.xticks([0,200,400,600],['','','',''])\n",
    "        else:\n",
    "            #plt.ticklabel_format(style='sci', axis='x', scilimits=(0,0))\n",
    "            plt.xticks([0,200,400,600],fontsize=fs)\n",
    "            plt.xlabel('[m$^{2}$ s$^{-1}$]',fontsize=fs)\n",
    "        plt.ylim(0,ylim2)\n",
    "        plt.yticks(yticks,yticklabs)\n",
    "        plt.grid(True,zorder=1)\n",
    "        if j == 0:\n",
    "            plt.title('K$_{h,v}$ (solid)\\nK$_{m,v}$ (dashed)',fontsize=fs2)\n",
    "            \n",
    "            \n",
    "            \n",
    "            \n",
    "        ax7 = plt.subplot(3,7,sp_loc[6][j])\n",
    "        plt_varx = mean_var_arr[i,time_arr[j],7,z1:z2]/mean_var_arr[i,time_arr[j],8,z1:z2]\n",
    "        #plt_vary = mean_var_arr[i,time_arr[j],-1,z1:z2]\n",
    "        plt_vary = theta_arr[i,1,time_arr[j],z1:z2]\n",
    "        plt.plot(plt_varx,plt_vary,c=colors[i],lw=lw,zorder=2,label=labels[i]) # KH/KM\n",
    "        #plt.scatter(plt_varx[::8],plt_vary[::8],s=60,facecolor=colors[i],edgecolor='k',zorder=2)\n",
    "        plt.xlim(0,2.2)\n",
    "        if j < 2:\n",
    "            plt.xticks([0,1,2],['','',''])\n",
    "        else:\n",
    "            plt.ticklabel_format(style='sci', axis='x', scilimits=(0,0))\n",
    "            plt.xticks([0,1,2],fontsize=fs)\n",
    "            plt.xlabel('[-]',fontsize=fs)\n",
    "        plt.ylim(0,ylim2)\n",
    "        plt.yticks(yticks,yticklabs)\n",
    "        plt.ylabel(str(time_arr[j]+1)+' h',rotation=-90,fontsize=fs2,labelpad=20)\n",
    "        plt.grid(True,zorder=1)\n",
    "        if j == 0:\n",
    "            plt.title('K$_{h,v}$/K$_{m,v}$',fontsize=fs2)\n",
    "            \n",
    "        ax7.yaxis.set_label_position(\"right\")\n",
    "                    \n",
    "plt.savefig(savedir+'traj_vprofs.png',dpi=600,bbox_inches='tight')\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FIG S4: PANEL SHOWING VPROFS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "var_lab = ['rthblten','rqvblten','rqcblten','rqiblten','rqniblten','rublten','rvblten','exch_h','exch_m','el_pbl','qke','sh3d','h_diabatic','qv_diabatic','qc_diabatic','ph']\n",
    "# np.empty([len(files_tends),len(calipso_j1),len(var_lab),np.shape(var)[0]])\n",
    "labels = [\"Level 2.5\", \"Level 3\", \"Level 2.5\\nEDMF\", \"YSU\"]\n",
    "\n",
    "z1 = 0\n",
    "z2 = 120\n",
    "lw = 1.5\n",
    "fs = 12\n",
    "fs2 = 14\n",
    "colors = ['lightgreen','lightblue','darkgreen','darkblue']\n",
    "time_arr = [-13,-7,-1]\n",
    "sp_loc = [[1,3,5],[2,4,6]]\n",
    "ylim2 = 7000\n",
    "yticks = [0,1000,2000,3000,4000,5000,6000,7000]\n",
    "yticklabs = ['0','','2000','','4000','','6000','']\n",
    "yticklabs2 = ['','','','','','','','']\n",
    "\n",
    "hydromet_arr2 = copy.deepcopy(hydromet_arr)\n",
    "\n",
    "plt.figure(figsize=(6,9))\n",
    "for i in np.arange(4): # sims\n",
    "    for j in np.arange(3): # times\n",
    "        plt.subplot(3,2,sp_loc[0][j])\n",
    "        plt_varx = 1e-6*hydromet_arr2[i,time_arr[j],5,z1:z2]*theta_arr[i,2,time_arr[j],z1:z2] # cc-1\n",
    "        #plt_vary = mean_var_arr[i,time_arr[j],-1,z1:z2]\n",
    "        plt_vary = theta_arr[i,1,time_arr[j],z1:z2]\n",
    "        plt.plot(plt_varx,plt_vary,c=colors[i],lw=lw,zorder=1) # KH/KM\n",
    "        #plt.scatter(plt_varx[::8],plt_vary[::8],s=60,facecolor=colors[i],edgecolor='k',zorder=2)\n",
    "        plt.xlim(0,120)\n",
    "        if j < 2:\n",
    "            plt.xticks([0,50,100],['','',''])\n",
    "        else:\n",
    "            plt.xticks([0,50,100],fontsize=fs)\n",
    "            plt.xlabel('CCN [cm$^{-3}$]',fontsize=fs2)\n",
    "        plt.ylim(0,ylim2)\n",
    "        plt.yticks(yticks,yticklabs,fontsize=fs)\n",
    "        plt.ylabel('Height ASL [m]',fontsize=fs2)\n",
    "        plt.grid(True)\n",
    "        \n",
    "        \n",
    "        \n",
    "        ax3 = plt.subplot(3,2,sp_loc[1][j])\n",
    "        plt_varx = 1e-6*hydromet_arr2[i,time_arr[j],7,z1:z2]*theta_arr[i,2,time_arr[j],z1:z2] # cc-1\n",
    "        #plt_vary = mean_var_arr[i,time_arr[j],-1,z1:z2]\n",
    "        plt_vary = theta_arr[i,1,time_arr[j],z1:z2]\n",
    "        ax3.plot(plt_varx,plt_vary,c=colors[i],lw=lw,zorder=1,label=labels[i]) # KH/KM\n",
    "        #plt.scatter(plt_varx[::8],plt_vary[::8],s=60,facecolor=colors[i],edgecolor='k',zorder=2)\n",
    "        ax3.set_xlim(0,65)\n",
    "        if j < 2:\n",
    "            ax3.set_xticks([0,30,60],['','',''])\n",
    "        else:\n",
    "            #ax3.ticklabel_format(style='sci', axis='x', scilimits=(0,0))\n",
    "            plt.xticks([0,30,60],fontsize=fs)\n",
    "            ax3.set_xlabel('CDNC [cm$^{-3}$]',fontsize=fs2)\n",
    "        if j == 0:\n",
    "            plt.legend()\n",
    "        ax3.set_ylim(0,ylim2)\n",
    "        ax3.set_yticks(yticks,yticklabs2,fontsize=fs)\n",
    "        plt.grid(True)\n",
    "        \n",
    "        ax3.yaxis.set_label_position(\"right\")\n",
    "        ax3.set_ylabel(str(time_arr[j]+1)+' h',rotation=-90,fontsize=fs2,labelpad=20)\n",
    "        \n",
    "#plt.show()\n",
    "plt.savefig(savedir+'traj_vprofs_na_nc.png',dpi=600,bbox_inches='tight')\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FIG S5: PANEL SHOWING VPROFS - Ni"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/glade/derecho/scratch/tjuliano/tmp/ipykernel_67821/4090030003.py:23: MatplotlibDeprecationWarning: Auto-removal of overlapping axes is deprecated since 3.6 and will be removed two minor releases later; explicitly call ax.remove() as needed.\n",
      "  ax5 = plt.subplot(3,1,sp_loc[j])\n",
      "/glade/derecho/scratch/tjuliano/tmp/ipykernel_67821/4090030003.py:23: MatplotlibDeprecationWarning: Auto-removal of overlapping axes is deprecated since 3.6 and will be removed two minor releases later; explicitly call ax.remove() as needed.\n",
      "  ax5 = plt.subplot(3,1,sp_loc[j])\n",
      "/glade/derecho/scratch/tjuliano/tmp/ipykernel_67821/4090030003.py:23: MatplotlibDeprecationWarning: Auto-removal of overlapping axes is deprecated since 3.6 and will be removed two minor releases later; explicitly call ax.remove() as needed.\n",
      "  ax5 = plt.subplot(3,1,sp_loc[j])\n",
      "/glade/derecho/scratch/tjuliano/tmp/ipykernel_67821/4090030003.py:23: MatplotlibDeprecationWarning: Auto-removal of overlapping axes is deprecated since 3.6 and will be removed two minor releases later; explicitly call ax.remove() as needed.\n",
      "  ax5 = plt.subplot(3,1,sp_loc[j])\n",
      "/glade/derecho/scratch/tjuliano/tmp/ipykernel_67821/4090030003.py:23: MatplotlibDeprecationWarning: Auto-removal of overlapping axes is deprecated since 3.6 and will be removed two minor releases later; explicitly call ax.remove() as needed.\n",
      "  ax5 = plt.subplot(3,1,sp_loc[j])\n",
      "/glade/derecho/scratch/tjuliano/tmp/ipykernel_67821/4090030003.py:23: MatplotlibDeprecationWarning: Auto-removal of overlapping axes is deprecated since 3.6 and will be removed two minor releases later; explicitly call ax.remove() as needed.\n",
      "  ax5 = plt.subplot(3,1,sp_loc[j])\n",
      "/glade/derecho/scratch/tjuliano/tmp/ipykernel_67821/4090030003.py:23: MatplotlibDeprecationWarning: Auto-removal of overlapping axes is deprecated since 3.6 and will be removed two minor releases later; explicitly call ax.remove() as needed.\n",
      "  ax5 = plt.subplot(3,1,sp_loc[j])\n",
      "/glade/derecho/scratch/tjuliano/tmp/ipykernel_67821/4090030003.py:23: MatplotlibDeprecationWarning: Auto-removal of overlapping axes is deprecated since 3.6 and will be removed two minor releases later; explicitly call ax.remove() as needed.\n",
      "  ax5 = plt.subplot(3,1,sp_loc[j])\n",
      "/glade/derecho/scratch/tjuliano/tmp/ipykernel_67821/4090030003.py:23: MatplotlibDeprecationWarning: Auto-removal of overlapping axes is deprecated since 3.6 and will be removed two minor releases later; explicitly call ax.remove() as needed.\n",
      "  ax5 = plt.subplot(3,1,sp_loc[j])\n"
     ]
    }
   ],
   "source": [
    "var_lab = ['rthblten','rqvblten','rqcblten','rqiblten','rqniblten','rublten','rvblten','exch_h','exch_m','el_pbl','qke','sh3d','h_diabatic','qv_diabatic','qc_diabatic','ph']\n",
    "# np.empty([len(files_tends),len(calipso_j1),len(var_lab),np.shape(var)[0]])\n",
    "labels = [\"Level 2.5\", \"Level 3\", \"Level 2.5\\nEDMF\", \"YSU\"]\n",
    "\n",
    "z1 = 0\n",
    "z2 = 120\n",
    "lw = 1.5\n",
    "fs = 12\n",
    "fs2 = 14\n",
    "colors = ['lightgreen','lightblue','darkgreen','darkblue']\n",
    "time_arr = [-13,-7,-1]\n",
    "sp_loc = [1,2,3]\n",
    "ylim2 = 7000\n",
    "yticks = [0,1000,2000,3000,4000,5000,6000,7000]\n",
    "yticklabs = ['','','','','','','','']\n",
    "yticklabs2 = ['0','','2000','','4000','','6000','']\n",
    "\n",
    "hydromet_arr2 = copy.deepcopy(hydromet_arr)\n",
    "\n",
    "plt.figure(figsize=(3,9))\n",
    "for i in np.arange(4): # sims\n",
    "    for j in np.arange(3): # times\n",
    "        ax5 = plt.subplot(3,1,sp_loc[j])\n",
    "        plt_varx = 1e-3*hydromet_arr2[i,time_arr[j],8,z1:z2]*theta_arr[i,2,time_arr[j],z1:z2] # L-1\n",
    "        #plt_vary = mean_var_arr[i,time_arr[j],-1,z1:z2]\n",
    "        plt_vary = theta_arr[i,1,time_arr[j],z1:z2]\n",
    "        ax5.plot(plt_varx,plt_vary,c=colors[i],lw=lw,zorder=1,label=labels[i]) # KH/KM\n",
    "        #plt.scatter(plt_varx[::8],plt_vary[::8],s=60,facecolor=colors[i],edgecolor='k',zorder=2)\n",
    "        plt.xlim(0,500)\n",
    "        if j < 2:\n",
    "            plt.xticks([0,250,500],['','',''])\n",
    "        else:\n",
    "        #    plt.ticklabel_format(style='sci', axis='x', scilimits=(0,0))\n",
    "            plt.xticks([0,250,500],fontsize=fs)\n",
    "            plt.xlabel('N$_{i}$ [L$^{-1}$]',fontsize=fs2)\n",
    "        ax5.set_ylim(0,ylim2)\n",
    "        plt.yticks(yticks,yticklabs2,fontsize=fs)\n",
    "        plt.ylabel('Height ASL [m]',fontsize=fs2)\n",
    "        plt.grid(True)\n",
    "        if j == 0:\n",
    "            plt.title('Temperature [$^{\\circ}$C]',fontsize=fs2)\n",
    "            plt.legend()\n",
    "            \n",
    "        #ax5.yaxis.set_label_position(\"right\")\n",
    "        \n",
    "        ax3a = ax5.twinx()\n",
    "        ax3b = ax5.twiny()\n",
    "        plt_varx = theta_arr[i,0,time_arr[j],z1:z2] - 273.15\n",
    "        plt_vary = theta_arr[i,1,time_arr[j],z1:z2]\n",
    "        ax3b.plot(plt_varx,plt_vary,c='k',lw=lw,zorder=1)\n",
    "        ax3b.plot([-35,-35],[plt_vary[0],plt_vary[-1]],c='b',lw=lw,ls='--',zorder=1)\n",
    "        ax3b.plot([-38,-38],[plt_vary[0],plt_vary[-1]],c='magenta',lw=lw,ls='--',zorder=1)\n",
    "        if j == 0:\n",
    "            ax3b.set_xticks([-50,-40,-30,-20,-10,0])\n",
    "            ax3b.tick_params(axis='x', which='major', labelsize=fs)\n",
    "        else:\n",
    "        #    plt.ticklabel_format(style='sci', axis='x', scilimits=(0,0))\n",
    "            ax3b.set_xticks([-50,-40,-30,-20,-10,0],['','','','','',''])\n",
    "        ax3b.set_xlim(-50,0)\n",
    "        #ax3b.yaxis.set_label_position(\"right\")\n",
    "        \n",
    "        ax3a.tick_params(top=False, labeltop=False, left=False, labelleft=False, right=False, labelright=False, bottom=False, labelbottom=False)\n",
    "        ax3b.tick_params(top=True, labeltop=True, left=False, labelleft=False, right=False, labelright=False, bottom=False, labelbottom=False)\n",
    "        ax3a.set_ylabel(str(time_arr[j]+1)+' h',rotation=-90,fontsize=fs2,labelpad=20)\n",
    "        ax3b.set_ylabel(str(time_arr[j]+1)+' h',rotation=-90,fontsize=fs2,labelpad=20)\n",
    "\n",
    "plt.savefig(savedir+'traj_vprofs_qn.png',dpi=600,bbox_inches='tight')\n",
    "plt.close()\n",
    "#plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
